# -*- coding: UTF-8 -*-
import io
import os
import time
from random import randint
from urllib.request import urlopen

import cv2
import numpy as np
from PIL import Image

from opencv_tool.HogGetter import makeHogForDir, getHogFeature2

# 计算的特征纬度，和HOG算法有关系
DIMEN = 144


class ImageTarget:
    '''
    封装分析出来的目标图片和X轴起始点
    '''

    def __init__(self, targetImg, x):
        self.targetImg = targetImg
        self.x = x
        self.recResult = ""


def analyseImage(path):
    '''
    分析GIF图片Model
    :param path:
    :return:
    '''
    # im = Image.open(path)

    if path.startswith("http"):
        image_bytes = urlopen(path).read()
        data_stream = io.BytesIO(image_bytes)
        im = Image.open(data_stream)
    else:
        im = Image.open(path)

    results = {
        'size': im.size,
        'mode': 'full',
    }
    try:
        while True:
            if im.tile:
                tile = im.tile[0]
                update_region = tile[1]
                update_region_dimensions = update_region[2:]
                if update_region_dimensions != im.size:
                    results['mode'] = 'partial'
                    break
            im.seek(im.tell() + 1)
    except EOFError:
        pass
    return results


def processImage(path):
    '''
    返回GIF所有帧
    :param path:
    :return:
    '''
    mode = analyseImage(path)['mode']
    if path.startswith("http"):
        image_bytes = urlopen(path).read()
        data_stream = io.BytesIO(image_bytes)
        im = Image.open(data_stream)
    else:
        im = Image.open(path)

    i = 0
    p = im.getpalette()
    last_frame = im.convert('RGBA')

    all_frame = []
    try:
        while True:
            if not im.getpalette():
                im.putpalette(p)
            new_frame = Image.new('RGBA', im.size)
            if mode == 'partial':
                new_frame.paste(last_frame)
            new_frame.paste(im, (0, 0), im.convert('RGBA'))
            img = cv2.cvtColor(np.asarray(new_frame), cv2.COLOR_RGB2BGR)
            all_frame.append(img)
            i += 1
            last_frame = new_frame
            im.seek(im.tell() + 1)
    except EOFError:
        pass

    return all_frame


def dealImage(path):
    '''
    处理得到的图片，切分联通区域
    :param path:
    :return:
    '''
    frames = processImage(path)
    result = frames[0]
    for index in range(1, len(frames)):
        f = frames[index]
        result = np.vstack([result, f])

    white = [255, 255, 255]

    result = cv2.copyMakeBorder(result, 0, 0, 10, 0, cv2.BORDER_CONSTANT, value=white)
    resultCopy = result.copy()

    result = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)

    height = result.shape[0]
    width = result.shape[1]

    ret, binary = cv2.threshold(result, 10, 10, cv2.THRESH_OTSU)
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    for con in contours:
        box = cv2.boundingRect(con)  # x，y，w，h
        if box[2] != width and (box[3] > height * 1.0 / 12):
            cv2.rectangle(result, (box[0], box[1]), (box[0] + box[2], box[1] + box[3]), (0, 0, 255), thickness=1)

    return result


def dealImageAndSplit(path):
    frames = processImage(path)
    result = frames[0]
    for index in range(1, len(frames)):
        f = frames[index]
        result = np.vstack([result, f])

    white = [255, 255, 255]

    result = cv2.copyMakeBorder(result, 0, 0, 10, 0, cv2.BORDER_CONSTANT, value=white)

    resultCopy = result.copy()

    result = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)

    height = result.shape[0]
    width = result.shape[1]

    ret, binary = cv2.threshold(result, 10, 10, cv2.THRESH_OTSU)
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    data_dir = "data"

    if not os.path.exists(data_dir):
        os.mkdir(data_dir)

    for con in contours:
        box = cv2.boundingRect(con)  # x，y，w，h
        if box[2] != width and (box[3] > height * 1.0 / 12):
            # x，y，w，h
            # cropped = img[0:128, 0:512]  # 裁剪坐标为[y0:y1, x0:x1]
            # 开始的y坐标:结束的y坐标,开始x:结束的x
            target = result[box[1]:box[1] + box[3], box[0]:box[0] + box[2]]
            cv2.rectangle(resultCopy, (box[0], box[1]), (box[0] + box[2], box[1] + box[3]), (0, 0, 255), thickness=1)

            cv2.imwrite(str(data_dir + os.path.sep + str(time.time()) + str(randint(1, 100000)) + ".png"), target)

    return resultCopy


def splitImage(path):
    '''
    把验证码拆分出来
    :param path:
    :return:
    '''
    frames = processImage(path)
    result = frames[0]
    for index in range(1, len(frames)):
        f = frames[index]
        result = np.vstack([result, f])

    white = [255, 255, 255]

    result = cv2.copyMakeBorder(result, 0, 0, 10, 0, cv2.BORDER_CONSTANT, value=white)

    result = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)

    height = result.shape[0]
    width = result.shape[1]

    ret, binary = cv2.threshold(result, 100, 120, cv2.THRESH_OTSU)
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    targets = []
    for con in contours:
        box = cv2.boundingRect(con)  # x，y，w，h
        if box[2] != width and (box[3] > height * 1.0 / 12):
            # x，y，w，h
            # cropped = img[0:128, 0:512]  # 裁剪坐标为[y0:y1, x0:x1]
            # 开始的y坐标:结束的y坐标,开始x:结束的x
            target = result[box[1]:box[1] + box[3], box[0]:box[0] + box[2]]
            targets.append(ImageTarget(target, box[0]))

    return targets


def rec_validate(path):
    '''
    识别指定文件、路径的图片
    :param path:
    :return:
    '''

    images = cv2.imread(path)
    # # 初始化机器学习引擎
    vec, list = makeHogForDir("./raw_data", dimen=DIMEN)
    label_list = []
    labelMap = {}
    for l in list:
        fname = int(os.path.basename(l).split(".")[0])
        label_list.append(fname)
        labelMap[int(fname)] = l

    label_sample = np.array(label_list).astype(np.float32).reshape((len(label_list), 1))
    knn = cv2.ml.KNearest_create()
    knn.train(vec, cv2.ml.ROW_SAMPLE, label_sample)

    rec_result = []
    index = 0
    for img in images:
        newcomer = getHogFeature2(img.targetImg).reshape((1, DIMEN))
        _, results, neighbours, dist = knn.findNearest(newcomer, 1)
        this_result = str(os.path.basename(os.path.dirname(labelMap[int(results)])))
        img.recResult = this_result
        rec_result.append(img)
        index = index + 1

    def takeSecond(elem):
        return elem.x

    rec_result.sort(key=takeSecond, reverse=False)

    uniqSet = []
    for recset in rec_result:
        if not uniqSet.__contains__(recset.recResult):
            uniqSet.append(recset.recResult)
    return uniqSet


if __name__ == "__main__":
    # collectdata(20, 3)
    start = time.time()
    rets = rec_validate("./test/TgqPp.gif")
    # rets = test1("http://www.opene164.org.cn/Num_Sys/checkcode/markquery/gifcaptcha.html")
    for ret in rets:
        print(ret)
    end = time.time()
    print("cost:" + str(end - start))
